#include "BPlusTree.h"
 
#ifdef TEST_TREE
#define PRIVATE 
#else
#define PRIVATE static
#endif

PRIVATE KeyType Unavailable = INT_MIN;

/* PRIVATE function declaration start */

PRIVATE BPlusTree MallocNewNode(pager_t *pager, int isLeaf);
PRIVATE BPlusTree RecursiveInsert(BPlusTree T, KeyType Key, int i, BPlusTree Parent, pager_t *pager, record_t *record, int *flag);
PRIVATE Position InsertElement(int isKey, Position Parent, Position X, KeyType Key, int i, int j, record_t *record);
PRIVATE Position FindSibling(Position Parent, int i);
PRIVATE Position FindSiblingKeyNum_M_2(Position Parent, int i, int *j);
PRIVATE Position FindMostLeft(Position P);
PRIVATE Position FindMostRight(Position P);
PRIVATE Position RemoveElement(int isKey, Position Parent, Position X, int i, int j);
PRIVATE Position MoveElement(Position Src, Position Dst, Position Parent, int i, int n);
PRIVATE BPlusTree SplitNode(Position Parent, Position X, int i, pager_t *pager);
PRIVATE Position MergeNode(Position Parent, Position X, Position S, int i);
PRIVATE BPlusTree RecursiveRemove(BPlusTree T, KeyType Key, int i, BPlusTree Parent);
PRIVATE void RecursiveTravel(BPlusTree T, int Level);
PRIVATE void MoveRecord(Position src, Position dest, KeyType src_idx);
PRIVATE void checkTree(BPlusTree tree);

/* PRIVATE function declaration end */
 
BPlusTree InitTree(pager_t *pager)
{
    BPlusTree T;    

    if (M < (3)) {
        fprintf(stderr, "Degree must be greater than 3!\n");
        exit(EXIT_FAILURE);
    }

    T = MallocNewNode(pager, 1);    /* root node */
    
    return T;
}

BPlusTree InsertToTree(BPlusTree T, KeyType Key, pager_t *pager, record_t *record, int *flag)
{
    return RecursiveInsert(T, Key, 0, NULL, pager, record, flag);
}

BPlusTree RemoveFromTree(BPlusTree T, KeyType Key)
{
    BPlusTree t = RecursiveRemove(T, Key, 0, NULL);
    checkTree(t);   /* stupid way */
    return t;
}
 
BPlusTree DestroyTree(BPlusTree T)
{
    int i;

    if (T != NULL) {
        i = 0;
        while (i < T->KeyNum + 1) {
            DestroyTree(T->Children[i]);
            i++;
        }
        // printf("Destroy:(");
        // j = 0;
        // while (j < T->KeyNum)   /* T->Key[i] != Unavailable */
        //     printf("%d:",T->Key[j++]);
        // printf(") ");

        free(T);
        T = NULL;
    }

    return T;
}
 
void Travel(BPlusTree T)
{
    RecursiveTravel(T, 0);
    printf("\n");
}
 
void TravelData(BPlusTree T)
{
    Position Tmp;
    int i;
	int count = 0;

    if (T == NULL)
        return ;
    printf("All Data:");

    Tmp = T;

    while (Tmp->Children[0] != NULL)
        Tmp = Tmp->Children[0];

    /* first leaf node */
    while (Tmp != NULL) {
        i = 0;
        while (i < Tmp->KeyNum) {
            printf(" %d",Tmp->Key[i++]);
			count++;
		}
        Tmp = Tmp->Next;
    }
    printf(" ----> Data Count: %d\n", count);
}

Position SearchFromTree(BPlusTree T, KeyType key)
{
    BPlusTree node = T;

    while (node != NULL) {
        int i = 0;
        while (i < node->KeyNum && key > node->Key[i])
            i++;

        if (node->Children[0] == NULL) {    /* leaf node */
            if (i < node->KeyNum && key == node->Key[i]) {
                return node;
            } else {
                return NULL;
            }
        }

        if (node->Key[i] == key)
            node = node->Children[i];
        else 
            node = node->Children[i - 1];
    }

    return NULL;
}

/* PRIVATE function implemention start */

/* generate and initialize a new node */
PRIVATE BPlusTree MallocNewNode(pager_t *pager, int isLeaf)
{
    BPlusTree NewNode;
    int idx;

    NewNode = malloc(sizeof(struct BPlusNode));
    if (NewNode == NULL) {
        fprintf(stderr, "malloc new node error\n");
        exit(EXIT_FAILURE);
    }
    
    idx = 0;
    while (idx < M + 1){
        NewNode->Key[idx] = Unavailable;
        NewNode->Children[idx] = NULL;
        idx++;
    }
    NewNode->Next = NULL;
    NewNode->KeyNum = 0;
    if (pager != NULL) {
        if (isLeaf) {
            if (getPage(pager, &(NewNode->page)) < 0) {
                fprintf(stderr, "get page in MallocNewNode error\n");
                destroyPager(pager);
                exit(EXIT_FAILURE);
            }
        } else {
            cleanPage(&(NewNode->page));
        }
    } else {
        NewNode->page.mem = NULL;
    }
    
    return NewNode;
}
 
PRIVATE Position FindMostLeft(Position P)
{
    Position Tmp;
    
    Tmp = P;
    
    while (Tmp != NULL && Tmp->Children[0] != NULL) {
        Tmp = Tmp->Children[0];
    }
    return Tmp;
}
 
PRIVATE Position FindMostRight(Position P)
{
    Position Tmp;
    
    Tmp = P;
    
    while (Tmp != NULL && Tmp->Children[Tmp->KeyNum - 1] != NULL) {
        Tmp = Tmp->Children[Tmp->KeyNum - 1];
    }
    return Tmp;
}
 
/*
 * Find a sibling node of the current node that has not yet reached
 * the maximum number of keys. If such a sibling is found, return it.
 * Otherwise, return NULL.
 * 
 * Parent->Children[i] points at the current node.
 */
PRIVATE Position FindSibling(Position Parent, int i)
{
    Position Sibling;
    int Limit;
    
    Limit = M;
    Sibling = NULL;

    if (i == 0) {   /* the current node is the first child */
        if (Parent->Children[1]->KeyNum < Limit)
            Sibling = Parent->Children[1];
    } else if (Parent->Children[i - 1]->KeyNum < Limit) {
        Sibling = Parent->Children[i - 1];
    } else if (i + 1 < Parent->KeyNum && Parent->Children[i + 1]->KeyNum < Limit){
        Sibling = Parent->Children[i + 1];
    }
    
    return Sibling;
}
 
/*
 * Find a sibling node of the current node that has more than M/2 keys.
 * If such a sibling is found, return it. Otherwise, return NULL.
 * 
 * Parent->Children[i] points at the current node.
 * j store the index of the sibling if found.
 */
PRIVATE Position FindSiblingKeyNum_M_2(Position Parent, int i, int *j)
{
    int Limit;
    Position Sibling;
    Sibling = NULL;
    
    Limit = LIMIT_M_2;
    
    if (i == 0) {       /* the current node is the first child */
        if (Parent->Children[1]->KeyNum > Limit) {
            Sibling = Parent->Children[1];
            *j = 1;
        }
    } else {
        if (Parent->Children[i - 1]->KeyNum > Limit){
            Sibling = Parent->Children[i - 1];
            *j = i - 1;
        }
        else if (i + 1 < Parent->KeyNum && Parent->Children[i + 1]->KeyNum > Limit){
            Sibling = Parent->Children[i + 1];
            *j = i + 1;
        }
    }
    
    return Sibling;
}
 
/* 
 * <<< parameters description >>>
 * 1. isKey = 1: Insert a key into node X
 *   (1) i ---> the index of node X in its parent (i.e., Parent->Children[i] points to X)
 *   (2) j ---> the index in node X where the key is to be inserted
 * 
 * 2. isKey = 0: Insert node X into Parent
 *   (1) i ---> the index in Parent where node X is to be inserted (i.e., Parent->Children[i] should point to X)
 *   (2) Key and j are not used when inserting node X
 */
PRIVATE Position InsertElement(int isKey, Position Parent, Position X, KeyType Key, int i, int j, record_t *record)
{
    int k;

    if (isKey) {        /* insert a key */
        k = X->KeyNum - 1;

        /* shift keys to the right to make space for the  new key */
        while (k >= j) {
            X->Key[k + 1] = X->Key[k];
            k--;
        }
        
        /* insert the new key at the correct position */
        X->Key[j] = Key;
        if (X->Children[0] == NULL && X->page.mem != NULL)
            insertRecordToPage(&(X->page), record); 
        
        /* update the parent's key (if it exists) */
        if (Parent != NULL)
            Parent->Key[i] = X->Key[0];
        
        X->KeyNum++;
    } else {            /* insert a node */
        /* connect the leaf node through next pointer if needed */
        if (X->Children[0] == NULL) {   /* X is the leaf node */
            if (i > 0) {
                X->Next = Parent->Children[i - 1]->Next;
                Parent->Children[i - 1]->Next = X;
            } else {
                X->Next = Parent->Children[0];
            }
        }
        
        k = Parent->KeyNum - 1;

        /* shift keys and children in the parent node to make space for the new node */
        while (k >= i) {
            Parent->Children[k + 1] = Parent->Children[k];
            Parent->Key[k + 1] = Parent->Key[k];
            k--;
        }

        /* insert the node's first key into the parent node */
        Parent->Key[i] = X->Key[0];
        Parent->Children[i] = X;
        
        Parent->KeyNum++;
    }

    return X;
}
 
PRIVATE Position RemoveElement(int isKey, Position Parent, Position X, int i, int j)
{
    int k, Limit;
    
    if (isKey) {        /* delete a key */
        Limit = X->KeyNum;
        if (X->Children[0] == NULL && X->page.mem != NULL)
            deleteRecordfromPage(&(X->page), X->Key[j]);
        k = j + 1;
        while (k < Limit) {
            X->Key[k - 1] = X->Key[k];
            k++;
        }
        
        X->Key[X->KeyNum - 1] = Unavailable;
        
        if (Parent != NULL)
            Parent->Key[i] = X->Key[0];
        
        X->KeyNum--;
    } else {        /* delete a node */

        /* modify the connect (next pointer) of leaf nodes */
        if (X->Children[0] == NULL && i > 0) {
            Parent->Children[i - 1]->Next = Parent->Children[i + 1];
        }
        Limit = Parent->KeyNum;
        k = i + 1;
        while (k < Limit) {
            Parent->Children[k - 1] = Parent->Children[k];
            Parent->Key[k - 1] = Parent->Key[k];
            k++;
        }
        
        Parent->Children[Parent->KeyNum - 1] = NULL;
        Parent->Key[Parent->KeyNum - 1] = Unavailable;
        
        Parent->KeyNum--;
        
    }
    return X;
}
 
/*
 * Src and Dst is the adjacent nodes.
 * i is the index that Src in Parent.
 * Move n elements from Src to Dst.
 */
PRIVATE Position MoveElement(Position Src, Position Dst, Position Parent, int i, int n)
{
    KeyType TmpKey;
    Position Child;
    int j,SrcInFront;
    
    SrcInFront = 0;
    
    if (Src->Key[0] < Dst->Key[0])
        SrcInFront = 1;
    
    j = 0;

    if (SrcInFront) {    /* Src comes before Dst */
        if (Src->Children[0] != NULL) {     /* internal node */
            while (j < n) {
                Child = Src->Children[Src->KeyNum - 1];
                RemoveElement(0, Src, Child, Src->KeyNum - 1, Unavailable);
                InsertElement(0, Dst, Child, Unavailable, 0, Unavailable, NULL);
                j++;
            }
        } else {    /* leaf node */
            while (j < n) {
                TmpKey = Src->Key[Src->KeyNum -1];
                /* store the record, which is needed to be moved */
                if (Src->page.mem != NULL) {
                    record_t record;
                    record.table = PAGE_GET((Src->page).type);
                    record.index = TmpKey;
                    char *start = getRecordFromPage(&(Src->page), record.index);
                    deserialize(start, &record);

                    RemoveElement(1, Parent, Src, i, Src->KeyNum - 1);
                    InsertElement(1, Parent, Dst, TmpKey, i + 1, 0, &record);
                } else {
                    RemoveElement(1, Parent, Src, i, Src->KeyNum - 1);
                    InsertElement(1, Parent, Dst, TmpKey, i + 1, 0, NULL);
                }
                j++;
            }
            
        }
        
        Parent->Key[i + 1] = Dst->Key[0];

        /* adjust the connect among leaf nodes */
        if (Src->KeyNum > 0)
            FindMostRight(Src)->Next = FindMostLeft(Dst);
    } else {    /* Dst comes before Src */
        if (Src->Children[0] != NULL) {     /* internal node */
            while (j < n) {
                Child = Src->Children[0];
                RemoveElement(0, Src, Child, 0, Unavailable);
                InsertElement(0, Dst, Child, Unavailable, Dst->KeyNum, Unavailable, NULL);
                j++;
            }
            
        } else {        /* leaf node */
            while (j < n) {
                TmpKey = Src->Key[0];
                /* store the record, which is needed to be moved */
                if (Src->page.mem != NULL) {
                    record_t record;
                    record.table = PAGE_GET((Src->page).type);
                    record.index = TmpKey;
                    char *start = getRecordFromPage(&(Src->page), record.index);
                    deserialize(start, &record);

                    RemoveElement(1, Parent, Src, i, 0);
                    InsertElement(1, Parent, Dst, TmpKey, i - 1, Dst->KeyNum, &record);
                } else {
                    RemoveElement(1, Parent, Src, i, 0);
                    InsertElement(1, Parent, Dst, TmpKey, i - 1, Dst->KeyNum, NULL);
                }
                j++;
            }
        }
        
        Parent->Key[i] = Src->Key[0];
        if (Src->KeyNum > 0)
            FindMostRight(Dst)->Next = FindMostLeft(Src);
    }
    
    return Parent;
}
 
PRIVATE BPlusTree SplitNode(Position Parent, Position X, int i, pager_t *pager)
{
    int j, k, Limit;
    Position NewNode;
    
    if (X->Children[0] == NULL) {
        NewNode = MallocNewNode(pager, 1);
    } else {
        NewNode = MallocNewNode(pager, 0);
    }
    
    k = 0;
    j = X->KeyNum / 2;
    Limit = X->KeyNum;
    while (j < Limit) {
        if (X->Children[0] != NULL) {   /* internal node (need to move children pointer to new node) */
            NewNode->Children[k] = X->Children[j];
            X->Children[j] = NULL;
        }
        /* move the key to new node */
        if (X->Children[0] == NULL && X->page.mem != NULL) {   /* leaf node (need to update its mapping page */
            MoveRecord(X, NewNode, X->Key[j]);
        }
        NewNode->Key[k] = X->Key[j];
        X->Key[j] = Unavailable;
        NewNode->KeyNum++; X->KeyNum--;
        j++; k++;
    }
    
    if (Parent != NULL)
        InsertElement(0, Parent, NewNode, Unavailable, i + 1, Unavailable, NULL);
    else{
        /* X is root (create a new root and return) */
        Parent = MallocNewNode(pager, 0);
        InsertElement(0, Parent, X, Unavailable, 0, Unavailable, NULL);
        InsertElement(0, Parent, NewNode, Unavailable, 1, Unavailable, NULL);
        
        return Parent;
    }
    
    return X;
}
 
/*
 * merge node
 * X has less than M/2 keys
 * S has greater than or equal M/2 keys
 */
PRIVATE Position MergeNode(Position Parent, Position X, Position S, int i)
{
    int Limit;
    
    if (S->KeyNum > LIMIT_M_2) {
        /* move one element from S to X */
        MoveElement(S, X, Parent, i, 1);
    } else {
        /* move all elements from X to S and free X */
        Limit = X->KeyNum;
        MoveElement(X, S, Parent, i, Limit);
        RemoveElement(0, Parent, X, i, Unavailable);
        
        free(X);
        X = NULL;
    }
    
    return Parent;
}
 
/* 
 * <<< parameters description >>>
 * (1) T ---> the current node where the insertion is performed
 * (2) key ---> the key to be inserted into the tree
 * (3) i ---> the index of the current node T in its parent node
 * (4) Parent ---> the parent node of T
 */
PRIVATE BPlusTree RecursiveInsert(BPlusTree T, KeyType Key, int i, BPlusTree Parent, pager_t *pager, record_t *record, int *flag)
{
    int j, Limit;
    Position Sibling;
    
    /* find the correct branch */
    j = 0;
    while (j < T->KeyNum && Key >= T->Key[j]) {
        if (Key == T->Key[j]) {   /* do not insert duplicate values */
            if (flag != NULL) *flag = 1;
            return T;
		}
        j++;
    }

    /*
    * j != 0 indicates we're not at the first key, so no adjustment is needed if it's the first key.
    * T->Children[0] != NULL checks if the node has children, meaning it's an internal node.
    * If the key to be inserted is less than T->key[j], the next recursive step should follow the left subtree,
    * i.e., we need to consider the child corresponding to T->key[j - 1].
    * Therefore, we decrement j by 1 to adjust the insertion position and move to the left child in the next recursion.
    */
    if (j != 0 && T->Children[0] != NULL) j--;
    
    if (T->Children[0] == NULL)     /* leaf node */
        T = InsertElement(1, Parent, T, Key, i, j, record);
    else    /* internal node */
        T->Children[j] = RecursiveInsert(T->Children[j], Key, j, T, pager, record, flag);
    
    /* adjust node if necessary */
    Limit = M;      /*maximum number of keys allowed in a node */

    /* need to split node */
    if (T->KeyNum > Limit) {
        if (Parent == NULL) {       /* root node */
            T = SplitNode(Parent, T, i, pager);
        } else {                    /* non-root node */
            Sibling = FindSibling(Parent, i);
            if (Sibling != NULL) {  /* find the sibling of the current node in the parent */
                MoveElement(T, Sibling, Parent, i, 1);  /* move an element (key or child) from T to its sibling */
            } else {                /* split the node if no sibling is found */
                T = SplitNode(Parent, T, i, pager);
            }
        }
    }
    
    if (Parent != NULL)
        Parent->Key[i] = T->Key[0];
    
    return T;
}
 
/* 
 * <<< parameters description >>>
 * T: the current node 
 * Key: the key that is being deleted from the tree
 * i: the index of the node T within its parent node's children
 * Parent: the parent node of T
 */
PRIVATE BPlusTree RecursiveRemove(BPlusTree T, KeyType Key, int i, BPlusTree Parent)
{
    int j, NeedAdjust;
    Position Sibling, Tmp;
    
    Sibling = NULL;
    
    /* find the correct branch */
    j = 0;
    while (j < T->KeyNum && Key >= T->Key[j]) {
        if (Key == T->Key[j])
            break;
        j++;
    }
    
    if (T->Children[0] == NULL) {
        /* Key not found */
        if (Key != T->Key[j] || j == T->KeyNum)
            return T;
    } else
        if (j == T->KeyNum || Key < T->Key[j]) j--;
    
    
    
    /* leaf node */
    if (T->Children[0] == NULL) {
        T = RemoveElement(1, Parent, T, i, j);
    } else {
        T->Children[j] = RecursiveRemove(T->Children[j], Key, j, T);
    }
    
    NeedAdjust = 0;
    /* root node, or leaf node, or number of children is between 2 and M */
    if (Parent == NULL && T->Children[0] != NULL && T->KeyNum < 2)
        NeedAdjust = 1;
    /* non-root internal nodes should have between ceil(M/2) and M children */
    else if (Parent != NULL && T->Children[0] != NULL && T->KeyNum < LIMIT_M_2)
        NeedAdjust = 1;
    /* (non-root) leaf nodes should also have between ceil(M/2) and M keys */
    else if (Parent != NULL && T->Children[0] == NULL && T->KeyNum < LIMIT_M_2)
        NeedAdjust = 1;
    
    /* adjust the node is necessary */
    if (NeedAdjust) {
        if (Parent == NULL) {   /* root node */
            /* 
             * the first child of root node is promoted to be the new root,
             * thereby reducing the height of the height of the tree
             */
            if(T->Children[0] != NULL && T->KeyNum < 2) {
                Tmp = T;
                T = T->Children[0];
                free(Tmp);
                return T;
            }
        } else {    /* non-root node */
            /* find a sibling with more than M/2 keys */
            Sibling = FindSiblingKeyNum_M_2(Parent, i,&j);
            if (Sibling != NULL) {
                MoveElement(Sibling, T, Parent, j, 1);
            } else {
                /* merge with sibling if no sibling with enough keys */
                if (i == 0)
                    Sibling = Parent->Children[1];
                else
                    Sibling = Parent->Children[i - 1];
                
                Parent = MergeNode(Parent, T, Sibling, i);
                T = Parent->Children[i];
            }
        }
    }

    return T;
}
 
PRIVATE void RecursiveTravel(BPlusTree T, int Level)
{
    int i;

    if (T != NULL){
        for (int i = 0; i <= Level; i++)
            printf("  ");
        printf("[Level:%d]-->",Level);
        printf("(");
        i = 0;
        while (i < T->KeyNum)/*  T->Key[i] != Unavailable*/
            printf("%d:",T->Key[i++]);
        printf(")\n");
        
        Level++;
        
        i = 0;
        while (i <= T->KeyNum) {
            RecursiveTravel(T->Children[i], Level);
            i++;
        }
    }
}

PRIVATE void MoveRecord(Position src, Position dest, KeyType src_idx)
{
    if (src->page.mem == NULL || dest->page.mem == NULL)
        return;

    record_t record;
    record.table = PAGE_GET((src->page).type);
    record.index = src_idx;

    char *start = getRecordFromPage(&(src->page), src_idx);
    deserialize(start, &record);
    insertRecordToPage(&(dest->page), &record);
    deleteRecordfromPage(&(src->page), src_idx);
}

PRIVATE void checkTree(BPlusTree tree)
{
    if (tree == NULL || tree->Children[0] == NULL) return;

    for (int i = 0; i < tree->KeyNum; i++) {
        checkTree(tree->Children[i]);
        tree->Key[i] = tree->Children[i]->Key[0];
    }
}

/* PRIVATE function implemention end */
